#include "../utils/load_region.hpp"
#include "../utils/load_point.hpp"
#include "../utils/load_element.hpp"
#include "../utils/apply_constraint_solve.hpp"
#include "../utils/L2error.hpp"
#include "../include/Grid.hpp"
#include "../../input/parameters.hpp"
#include <array>
#include <Eigen/Dense>
#include <Eigen/Sparse>
#include <Eigen/IterativeLinearSolvers>

int main()
{
    std::cout << "running program" << std::endl;
    Region ref = load_region("triangle");
    std::vector<Point> p = load_point();
    std::cerr << "points vector has been constructed!" << std::endl;
    std::vector<Element> e = load_element(p);
    std::cerr << "elements vector has been constructed!" << std::endl;
    Grid g(p, e);
    unsigned long mat_size = p.size();
    std::cerr << "the size of matrix is " << mat_size << std::endl;
    Eigen::SparseMatrix<double> A(mat_size, mat_size);
    Eigen::VectorXd b = Eigen::VectorXd::Zero(mat_size);
    Eigen::VectorXd x = Eigen::VectorXd::Zero(mat_size);

    std::function<double(double, double)> f = [](double x, double y)
    {
        return RHS(x,y);
    };
    std::function<double(double, double)> gd = [](double x, double y)
    {
        return dirichlet(x,y);
    };

    std::cout << "begin to assemble matrix and vector" << std::endl; 
    if(element_shape == 3)
    {
        if(degree_of_basis_function == 1)
        {
//            #pragma omp parallel for num_threads(4)
            for(auto eiter=g.elements.begin(); eiter<g.elements.end(); eiter++)
            {
//                std::cerr << "calculate " << (*eiter);
                double J = Jacobi((*eiter));
//                std::cerr << "det(Jacobi)=" << J << std::endl;
                unsigned long p0_idx = id((*eiter)[0]);
                unsigned long p1_idx = id((*eiter)[1]);
                unsigned long p2_idx = id((*eiter)[2]);
                double b2 = ((*eiter)[2])[1] - ((*eiter)[0])[1];
                double b3 = ((*eiter)[0])[1] - ((*eiter)[1])[1];
                double c2 = ((*eiter)[0])[0] - ((*eiter)[2])[0];
                double c3 = ((*eiter)[1])[0] - ((*eiter)[0])[0];
                //prepare integrator for calculate stiff matrix
                std::function<double(double, double)> h00 = [b2, b3, c2, c3, &ref, J](double x, double y)
                {
                   return ((b2 * ref.basis_gradient[0][0](x,y) + b3 * ref.basis_gradient[0][1](x,y)) * \
                          (b2 * ref.basis_gradient[0][0](x,y) + b3 * ref.basis_gradient[0][1](x,y)) + \
                          (c2 * ref.basis_gradient[0][0](x,y) + c3 * ref.basis_gradient[0][1](x,y)) * \
                          (c2 * ref.basis_gradient[0][0](x,y) + c3 * ref.basis_gradient[0][1](x,y))) / J; 
                };
                std::function<double(double, double)> h01 = [b2, b3, c2, c3, &ref, J](double x, double y)
                {
                   return ((b2 * ref.basis_gradient[0][0](x,y) + b3 * ref.basis_gradient[0][1](x,y)) * \
                          (b2 * ref.basis_gradient[1][0](x,y) + b3 * ref.basis_gradient[1][1](x,y)) + \
                          (c2 * ref.basis_gradient[0][0](x,y) + c3 * ref.basis_gradient[0][1](x,y)) * \
                          (c2 * ref.basis_gradient[1][0](x,y) + c3 * ref.basis_gradient[1][1](x,y))) / J; 
                };
                std::function<double(double, double)> h02 = [b2, b3, c2, c3, &ref, J](double x, double y)
                {
                   return ((b2 * ref.basis_gradient[0][0](x,y) + b3 * ref.basis_gradient[0][1](x,y)) * \
                          (b2 * ref.basis_gradient[2][0](x,y) + b3 * ref.basis_gradient[2][1](x,y)) + \
                          (c2 * ref.basis_gradient[0][0](x,y) + c3 * ref.basis_gradient[0][1](x,y)) * \
                          (c2 * ref.basis_gradient[2][0](x,y) + c3 * ref.basis_gradient[2][1](x,y))) / J; 
                };
                std::function<double(double, double)> h11 = [b2, b3, c2, c3, &ref, J](double x, double y)
                {
                   return ((b2 * ref.basis_gradient[1][0](x,y) + b3 * ref.basis_gradient[1][1](x,y)) * \
                          (b2 * ref.basis_gradient[1][0](x,y) + b3 * ref.basis_gradient[1][1](x,y)) + \
                          (c2 * ref.basis_gradient[1][0](x,y) + c3 * ref.basis_gradient[1][1](x,y)) * \
                          (c2 * ref.basis_gradient[1][0](x,y) + c3 * ref.basis_gradient[1][1](x,y))) / J; 
                };
                std::function<double(double, double)> h12 = [b2, b3, c2, c3, &ref, J](double x, double y)
                {
                   return ((b2 * ref.basis_gradient[1][0](x,y) + b3 * ref.basis_gradient[1][1](x,y)) * \
                          (b2 * ref.basis_gradient[2][0](x,y) + b3 * ref.basis_gradient[2][1](x,y)) + \
                          (c2 * ref.basis_gradient[1][0](x,y) + c3 * ref.basis_gradient[1][1](x,y)) * \
                          (c2 * ref.basis_gradient[2][0](x,y) + c3 * ref.basis_gradient[2][1](x,y))) / J; 
                };
                std::function<double(double, double)> h22 = [b2, b3, c2, c3, &ref, J](double x, double y)
                {
                   return ((b2 * ref.basis_gradient[2][0](x,y) + b3 * ref.basis_gradient[2][1](x,y)) * \
                          (b2 * ref.basis_gradient[2][0](x,y) + b3 * ref.basis_gradient[2][1](x,y)) + \
                          (c2 * ref.basis_gradient[2][0](x,y) + c3 * ref.basis_gradient[2][1](x,y)) * \
                          (c2 * ref.basis_gradient[2][0](x,y) + c3 * ref.basis_gradient[2][1](x,y))) / J; 
                };
                // calculate stiff matrix
                double A00 = ref.integral(h00);
                double A01 = ref.integral(h01);
                double A02 = ref.integral(h02);
                double A11 = ref.integral(h11);
                double A12 = ref.integral(h12);
                double A22 = ref.integral(h22);
                // std::cerr << "unit stiffness matrix" << std::endl;
                // std::cerr << A00 << " " << A01 << " " << A02 << std::endl;
                // std::cerr << A01 << " " << A11 << " " << A12 << std::endl;
                // std::cerr << A02 << " " << A12 << " " << A22 << std::endl;
                // assemble stiff matrix 
                A.coeffRef(p0_idx, p0_idx) += A00;
                A.coeffRef(p0_idx, p1_idx) += A01;
                A.coeffRef(p0_idx, p2_idx) += A02;
                A.coeffRef(p1_idx, p1_idx) += A11;
                A.coeffRef(p1_idx, p2_idx) += A12;
                A.coeffRef(p2_idx, p2_idx) += A22;
                A.coeffRef(p1_idx, p0_idx) += A01;
                A.coeffRef(p2_idx, p0_idx) += A02;
                A.coeffRef(p2_idx, p1_idx) += A12;
//                std::cout << "iteration=" << cnt << std::endl;
//                std::cout << "matrix A =" << std::endl;
//                std::cout << A << std::endl;

                double rhs_p0 = 0.0;
                double rhs_p1 = 0.0;
                double rhs_p2 = 0.0;
                for(auto niter = ref.nodes.begin(); niter<ref.nodes.end(); niter++)
                {
                    std::array<double, 2> global_node = (*eiter).l2g((*niter)[0], (*niter)[1]);
                    double temp = f(global_node[0], global_node[1]) * J * ((*niter).get_weight()) / 2.0;
                    rhs_p0 += ref.basis_function[0]((*niter)[0], (*niter)[1]) * temp;
                    rhs_p1 += ref.basis_function[1]((*niter)[0], (*niter)[1]) * temp;   
                    rhs_p2 += ref.basis_function[2]((*niter)[0], (*niter)[1]) * temp;
                }
                Eigen::VectorXd unitrhs = Eigen::VectorXd::Zero(mat_size);
                unitrhs[p0_idx] = rhs_p0;
                unitrhs[p1_idx] = rhs_p1;
                unitrhs[p2_idx] = rhs_p2;
                b += unitrhs;
            }
            std::cout << "solve Ax = b by simplicial LDLT method" << std::endl;
            x = apply_constraint_solve(A, b, g, gd);
        }
        else
        {
        //    std::cerr << "element_shape =" << element_shape << " and basis_function is P2" << std::endl; 
            for(auto eiter=g.elements.begin(); eiter<g.elements.end(); eiter++)
            {
        //        std::cerr << "calculating " << (*eiter);
                double J = Jacobi((*eiter));
                double b2 = ((*eiter)[2])[1] - ((*eiter)[0])[1];
                double b3 = ((*eiter)[0])[1] - ((*eiter)[1])[1];
                double c2 = ((*eiter)[0])[0] - ((*eiter)[2])[0];
                double c3 = ((*eiter)[1])[0] - ((*eiter)[0])[0];
        //        std::cerr << "enter the loop" << std::endl;
                for(int bf1=0; bf1<6; bf1++)
                {
                    unsigned long pi_idx = id((*eiter)[bf1]);
        //            std::cerr << "calculate matrix A for Point's id=" << pi_idx; 
                    for(int bf2=bf1; bf2<6; bf2++)
                    {
                        unsigned long pj_idx = id((*eiter)[bf2]);
        //                std::cerr << " and Point's id" << pj_idx << std::endl; 
                        std::function<double(double, double)> h = [bf1, bf2, b2, b3, c2, c3, &ref, J](double x, double y)
                        {
                            return ((b2 * ref.basis_gradient[bf1][0](x,y) + b3 * ref.basis_gradient[bf1][1](x,y)) * \
                                    (b2 * ref.basis_gradient[bf2][0](x,y) + b3 * ref.basis_gradient[bf2][1](x,y)) + \
                                    (c2 * ref.basis_gradient[bf1][0](x,y) + c3 * ref.basis_gradient[bf1][1](x,y)) * \
                                    (c2 * ref.basis_gradient[bf2][0](x,y) + c3 * ref.basis_gradient[bf2][1](x,y))) / J; 
                        };
                        
                        double Aij = ref.integral(h);
                        if(bf1 == bf2)
                        {
                            A.coeffRef(pi_idx, pj_idx) += Aij;
                        }
                        else
                        {
                            A.coeffRef(pi_idx, pj_idx) += Aij;
                            A.coeffRef(pj_idx, pi_idx) += Aij;
                        }
                    }
        //            std::cerr << "calculate vector b for Point " << pi_idx << std::endl;
                    double rhs_value = 0.0; 
                    for(auto niter = ref.nodes.begin(); niter<ref.nodes.end(); niter++)
                    {
                        std::array<double, 2> global_node = (*eiter).l2g((*niter)[0], (*niter)[1]);
                        double temp = f(global_node[0], global_node[1]) * J * ((*niter).get_weight()) / 2.0;
                        rhs_value += ref.basis_function[bf1]((*niter)[0], (*niter)[1]) * temp;
                    }
                    b[pi_idx] += rhs_value;
                }
           }
            std::cout << "solve Ax = b by simplicial LDLT method" << std::endl;
            x = apply_constraint_solve(A, b, g, gd);
        }
    }
    else
    {
        std::cerr << "Quadrilateral element is not ready yet!" << std::endl;
    }
     
    //L2 error
    double error = L2error(g, x, ref, gd);
    std::cerr << "length of element side and L2 error=" << error << std::endl; 
    // for(unsigned long idx=0; idx<mat_size; idx++)
    // {
        // diff[idx] = abs(diff[idx]/b[idx]);
    // }
    
//    /** print data for plot 
    std::ofstream fout;
    fout.open("../data/plot.txt");
    for(unsigned i=0; i<p.size(); i++)
    {
        fout << ((g.vertex[i])[0]) << " " << ((g.vertex[i])[1]) << " " << (x(i)) << std::endl;
    }
    fout.close();
//    **/
    
    
    
    
    // std::cout << "check error each coordinate using ConjugateGradient method" << std::endl;
    // Eigen::VectorXd diff(mat_size);
    // diff = A*x - b;
    // for(unsigned long row=0; row<mat_size; row++ )
    // {
        // std::cout << diff(row) << std::endl;
    // }

    return 0;
}
